#! /usr/bin/env python
# -*- coding: utf-8 -*-

# Copyright (c) JoinQuant Development Team
# Author: Huayong Kuang <kuanghuayong@joinquant.com>

import time
import socket
import logging
import threading
from concurrent.futures import ThreadPoolExecutor

from werkzeug.urls import uri_to_iri
from werkzeug.serving import reraise, InternalServerError
from werkzeug.serving import BaseWSGIServer, WSGIRequestHandler, run_simple

from flask import Flask


app = Flask(__name__)


class WSGIServer(BaseWSGIServer):

    def shutdown_request(self, request):
        """Called to shutdown and close an individual request."""
        is_thread_handler = getattr(self.RequestHandlerClass, "is_thread", False)
        if is_thread_handler:
            # 如果是多线程的处理器，则不关闭连接，由线程自己关闭
            return
        self.close_request()

    def close_request(self, request):
        """Called to clean up an individual request."""
        try:
            # explicitly shutdown.  socket.close() merely releases
            # the socket and waits for GC to perform the actual close.
            request.shutdown(socket.SHUT_WR)
        except OSError:
            pass  # some platforms may raise ENOTCONN here
        request.close()


class ThreadedRequestHandler(WSGIRequestHandler):

    is_thread = True  # 标记为是多线程请求处理器
    status = None  # 响应状态信息
    response_headers = None  # 响应体头部数据
    code = 0  # 响应状态码
    time_start = 0  # 请求开始处理的时间
    time_finish = 0  # 请求处理完成的时间
    headers_sent = False  # 标记响应头部数据是否发送
    response_length = 0  # 响应数据的大小
    result = None  # WSGI 应用的返回值
    content_length = 0  # 响应内容部分大小
    iri_path = None  # 国际化资源标识路径

    __log = logging.getLogger(__name__)

    pool1 = ThreadPoolExecutor(max_workers=5, thread_name_prefix="test")
    pool2 = ThreadPoolExecutor(max_workers=1)

    def __init__(self, request, client_address, server):
        self.request = request
        self.client_address = client_address
        self.server = server
        self.setup()
        self.handle_in_thread()

    def __call__(self):
        try:
            self.handle()
        except Exception as e:
            self.server.handle_error(self.request, self.client_address)
            self._log("exception", str(e))
        finally:
            self.finish()

    def _log(self, type, message, *args):
        message = (message % args).rstrip()
        log = getattr(self.__log, type.lower())
        log("({}) -- {}".format(
            self._get_tread_name(), message
        ))

    def log(self, type, message, *args):
        message = '{} - "{}" {}'.format(
            self.address_string(),
            self._request_string(),
            message
        )
        self._log(type, message, *args)

    @staticmethod
    def _get_tread_name():
        return threading.current_thread().getName()

    @property
    def server_version(self):
        return "TestServer/0.1"

    def parse_request(self):
        ret = super().parse_request()
        self.iri_path = uri_to_iri(self.path)
        return ret

    def _write(self, data):
        assert isinstance(data, bytes), "applications must write bytes"
        size = self.wfile.write(data)
        self.wfile.flush()
        self.response_length += size
        return size

    def send_header_data(self):
        size = 0
        if self.headers.get("Expect", "").lower().strip() == "100-continue":
            size += self._write(b"HTTP/1.1 100 Continue\r\n\r\n")

        if self.request_version != "HTTP/0.9":
            hdr = "%s %s\r\n" % (self.protocol_version, self.status)
            size += self._write(hdr.encode("ascii"))

        environ = self.environ
        code = self.code
        header_keys = set()
        for key, value in self.response_headers:
            self.send_header(key, value)
            key = key.lower()
            header_keys.add(key)
        if not (
            "content-length" in header_keys
            or environ["REQUEST_METHOD"] == "HEAD"
            or code < 200
            or code in (204, 304)
        ):
            self.close_connection = True
            self.send_header("Connection", "close")
        if "server" not in header_keys:
            self.send_header("Server", self.version_string())
        if "date" not in header_keys:
            self.send_header("Date", self.date_time_string())

        if self.request_version != 'HTTP/0.9':
            self._headers_buffer.append(b"\r\n")
            if hasattr(self, '_headers_buffer'):
                size += self._write(b"".join(self._headers_buffer))
                self._headers_buffer = []

        self.headers_sent = True
        self.log("debug", "sent header data, size: %s", size)
        return size

    def write(self, data):
        if self.code in (304, 204) and data:
            raise AssertionError('The %s response must have no body' % self.code)

        if not self.headers_sent:
            if not self.status:
                raise AssertionError("The application did not call start_response()")
            self.send_header_data()
        return self._write(data)

    def start_response(self, status, response_headers, exc_info=None):
        if exc_info:
            try:
                if self.headers_sent:
                    reraise(*exc_info)
            finally:
                exc_info = None

        try:
            code, msg = status.split(None, 1)
        except ValueError:
            code, msg = status, ""
        if msg is None:
            msg = code in self.responses and self.responses[code][0] or ""

        self.status = status
        self.code = int(status.split(None, 1)[0])
        self.response_headers = response_headers

        return self.write

    def send_body_data(self, data=None):
        data = data or self.result
        if isinstance(data, bytes):
            size = self.write(data)
        elif isinstance(data, str):
            size = self.write(data.encode("utf-8"))
        else:
            size = sum([self.write(_d) for _d in data])

        if size == 0:
            self.write(b"")

        self.content_length = size
        self.log("debug", "sent body data, size: %s", size)
        return size

    def execute_app(self, app=None):
        app = app or self.server.app
        self.result = app(self.environ, self.start_response)
        try:
            self.send_body_data()
        finally:
            if hasattr(self.result, "close"):
                self.result.close()
            self.result = None

    def run_wsgi(self):
        self.environ = environ = self.make_environ()

        try:
            self.execute_app()
        except (ConnectionError, socket.timeout) as e:
            self.connection_dropped(e, environ)
        except Exception:
            if self.server.passthrough_errors:
                raise
            from werkzeug.debug.tbtools import get_current_traceback

            traceback = get_current_traceback(ignore_system_exceptions=True)
            try:
                self.execute_app(InternalServerError())
            except Exception:
                pass
            self.server.log("error", "Error on request:\n%s", traceback.plaintext)

    def handle_one_request(self):
        """处理单个 HTTP 请求"""
        try:
            return self.run_wsgi()
        finally:
            self.time_finish = time.time()
            self.log_request(self.code, self.response_length)

    def handle_in_thread(self):
        self.time_start = time.time()
        self.raw_requestline = self.rfile.readline()
        if not self.raw_requestline:
            self.close_connection = 1
        elif self.parse_request():
            self.pool1.submit(self)
            self.__log.debug('submitted request "%s" to %s',
                             self._request_string(), self.pool1)
        else:
            self.__log.error('parse request "%s %s" failure',
                             self.request, self.client_address)

    def finish(self):
        super().finish()
        self.server.close_request()

    def _request_string(self):
        try:
            request = "{} {} {}".format(
                self.command,
                self.iri_path,
                self.request_version
            )
        except AttributeError:
            # path isn't set if the requestline was bad
            request = self.requestline
        return request

    def log_request(self, code="-", size="-"):
        code = str(code)
        if size == 0:
            size = "-"
        if self.time_finish:
            delta = '%.6f' % (self.time_finish - self.time_start)
        else:
            delta = '-'
        self.log("info", '%s %s %s', code, size, delta)


@app.route('/')
def hello_world():
    __import__("time").sleep(1)
    return 'Hello, World!'


@app.route("/<bar>")
def foo(bar):
    return "----- " + bar


if __name__ == "__main__":
    logging.basicConfig(level=logging.DEBUG,
                        format='%(asctime)s - %(levelname)s - %(message)s')
    server = WSGIServer("localhost", 8808, app, handler=ThreadedRequestHandler)
    print("serve on localhost:8808")
    server.serve_forever()
    # run_simple("localhost", 8808, app)
